#!/usr/bin/env python
#-*- coding:utf-8 -*-

"""
Trackers

todo:
    wrap KLT + gpu = http://cs.unc.edu/~ssinha/Research/GPU_KLT/

"""

from __future__ import division

from tools import *
from visualise import *
from drawing import *
from features import *

from kdtree import KDTree

from collections import deque

class FIFO(deque):
    def __init__(self, max_size=None):
        self.max_size = max_size if max_size is not None else Inf
        deque.__init__(self)

    def add(self, item):
        self.append(item)
        if len(self) > self.max_size:
            pop_item = self.popleft()
            return pop_item
        return None

    def is_full(self):
        return len(self) >= self.max_size

class TimeFIFO(object):
    def __init__(self, max_age=Inf):
        self.points = deque()
        self.times = deque()
        self.max_age = max_age

    def add(self, point, the_time):
        self.points.append(point)
        self.times.append(the_time)
        while 1:
            if not self.times:
                break
            if (the_time - self.times[0]) <= self.max_age: # oldest item still valid
                break

            del_point = self.points.popleft()
            del_the_time = self.times.popleft()

    def last_pos(self):
        return self.points[-1] if self.points else None

    def __bool__(self):
        return bool(self.points)

    def __array__(self):
        return asarray(self.points)

    def __iter__(self):
        return iter(self.points)

    def __len__(self):
        return len(self.times)

class TrackedPoint(object):
    def __init__(self, pos, the_time):
        self.positions = TimeFIFO(max_age=2)
        self.positions.add(pos, the_time)

    def update_pos(self, pos, the_time):
        self.positions.add(pos, the_time)

    @property
    def pos(self):
        return self.positions.last_pos()

    def draw(self, image):
        draw_polyline(image, points=list(self.positions)[-3:])

@memoize
def circular_offsets(max_x, max_y):
    db = shelve.open(db_filename)
    try:
        entry = db["circular_offsets"]
    except KeyError:
        entry = {}
    try:
        return entry[(max_x, max_y)]
    except KeyError:
        pass
    logger.debug("determining list of circular offsets for grid %d x %d" % (max_x, max_y))
    items = []
    for y in range(-max_y+1, max_y):
        for x in range(-max_x+1, max_x):
            items.append((x,y))
    items.sort(key=lambda (x,y): (x**2+y**2, x, y))
    items = asarray(items)

    entry[(max_x, max_y)] = items
    db["circular_offsets"] = entry
    db.sync()
    logger.debug("done")
    return items

_pool = None
def get_pool():
    global _pool
    if _pool == None:
        _pool = processing.Pool()
    return _pool

import processing

class Grid(object):
    def __init__(self, size):
        self.items = ones(size, dtype=int) * -1
        self.size = size

    def add(self, pos, item):
        self.items[int(pos[0]), int(pos[1])] = item

    def nearest_to_positions(self, positions):
        #~ return self._nearest_to_positions(positions)

        the_pool = get_pool()
        N = len(positions)
        results = []
        for positions_part in [positions[:N//2], positions[N//2:]]:
            result = the_pool.apply_async(self._nearest_to_positions, args=(positions_part,), kwds=dict(indices=circular_offsets(*self.size)))
            results.append(result)
        for result in results:
            result.wait()
        return hstack([result.get() for result in results])

    def _nearest_to_positions(self, positions, indices=None):
        X, Y = self.size
        if indices is None:
            indices = circular_offsets(*self.size)
        N = len(indices)
        P = len(positions)
        items = self.items
        results = zeros(len(positions), dtype=int)
        code = """
            int x, y, item;
            for (int pn=0; pn < P; pn++) {
                int x0 = positions(pn, 0);
                int y0 = positions(pn, 1);
                results(pn) = -1;

                for (int n=0; n < N; n++) {
                    x = x0 + indices(n, 0);
                    y = y0 + indices(n, 1);
                    if ((x < 0) || (X <= x))
                        continue;
                    if ((y < 0) || (Y <= y))
                        continue;
                    item = items(x, y);
                    if (item != -1) {
                        results(pn) = item;
                        break;
                    }
                }

            }
        """
        scipy.weave.inline(code,
            ["positions", "indices", "X", "Y", "N", "items", "results", "P"],
            type_converters=scipy.weave.converters.blitz,
            compiler = 'gcc')
        return [(r if r != -1 else None) for r in results]

class HashGrid(object):
    def __init__(self, size, max_distance=10):
        self.size = size
        self.block_size = asarray([max_distance // 2, max_distance // 2])
        self.items = {}

    def add(self, pos, item):
        x, y = map(int, pos // self.block_size)
        self.items.setdefault((x, y), []).append((pos, item))

    def nearest_to_positions(self, positions):
        results = []
        for pos in positions:
            x, y = map(int, pos // self.block_size)
            items = []
            for dx in [-1, 0, 1]:
                for dy in [-1, 0, 1]:
                    items.extend(self.items.get((x+dx, y+dy), []))

            if items:
                n = argmin([((pos-p)**2).sum() for p, item in items])
                p, item = items[n]
                results.append(item)
            else:
                results.append(None)

        return results

class Tracker(object):
    def __init__(self, image_size=None):
        self.prev_points = []
        self.image_size = image_size

    def __iter__(self):
        return iter(self.prev_points)

    def add_features(self, features, the_time):
        return self._add_features_grid(features, the_time)
        return self._add_features_kdtree(features, the_time)

    def _add_features_kdtree(self, features, the_time):

        if not features:
            return

        if not self.prev_points:
            self.prev_points = [TrackedPoint(f.pos, the_time) for f in features]
            return

        new_observations = [f.pos for f in features]
        new_observations_is_novel = ones(len(new_observations))
        tree = KDTree(new_observations)
        closest_d, closest_indices = tree.query([point.pos for point in self.prev_points])

        for point, n, d in zip(self.prev_points, closest_indices, closest_d):
            point.update_pos(new_observations[n], the_time)
            new_observations_is_novel[n] = 0

        for obs, is_novel in zip(new_observations, new_observations_is_novel):
            if is_novel:
                self.prev_points.append(TrackedPoint(obs, the_time))

    def _add_features_grid(self, features, the_time):

        if not features:
            return

        if not self.prev_points:
            self.prev_points = [TrackedPoint(f.pos, the_time) for f in features]
            return

        new_observations = [f.pos for f in features]
        new_observations_is_novel = ones(len(new_observations))
        new_points_grid = Grid(size=self.image_size)
        for p, n in zip(new_observations, count()):
            new_points_grid.add(p, n)

        old_observations = asarray([point.pos for point in self.prev_points])
        ns = new_points_grid.nearest_to_positions(old_observations )
        for n, point in zip(ns, self.prev_points):
            if n is None:
                continue
            try:
                point.update_pos(new_observations[n], the_time)
            except TypeError:
                logger.error(n)
                raise
            new_observations_is_novel[n] = 0

        for obs, is_novel in zip(new_observations, new_observations_is_novel):
            if is_novel:
                self.prev_points.append(TrackedPoint(obs, the_time))

    def add_frame(self, frame):
        image = frame["change"]
        image = frame["clean"]
        self.image_size = image.shape[:2]

        features = get_fast_corners(image, threshold=20)
        #~ features = get_mser_features(image)
        #~ features = get_surf_features(image)

        t = time.time()
        self.add_features(features, frame["the_time"])
        logger.debug("tracked features in %0.2fs" % (time.time() - t))

        image = frame["display"]

        for point in self.prev_points:
            point.draw(image)
        #~ for feature in features:
            #~ feature.draw(image)

def main():
    v = Viewer()

    filename = "http://146.232.169.181/mjpg/video.mjpg" # shared area
    filename = "../media/test.avi"
    filename = "../media/test.mpg"
    tracker = Tracker()
    for frame in Video(filename).frames():
        tracker.add_frame(frame)

        v.update_image(frame["display"])

        v.draw()
        if not v.alive:
            break

if __name__ == '__main__':
    prof(main)
    #~ main()
